from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
import os


def read_data(filename) -> str:
    path = filename
    
    with open(path, mode="r") as f:
        return f.read()


def get_chunks(filename=""):
    if not filename:
        filename = "samples/default.md"

    markdown = read_data(filename)

    headers_to_split_on = [
        ("#", "Heading 1"),
        ("##", "Heading 2"),
        ("###", "Heading 3"),
        ("####", "Heading 4"),    
    ]

    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on,
        strip_headers=False
    )
    md_header_splits = markdown_splitter.split_text(markdown)

    chunk_size = 500
    chunk_overlap = 100
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )

    return text_splitter.split_documents(md_header_splits)


def get_chunks_content_only(filename="") -> str:
    return [d.page_content for d in get_chunks(filename)]


if __name__ == "__main__":
    for d in get_chunks_content_only():
        print(f"{'=' * 30}\n{d}\n")
